package com.hdu.ordered;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.LongAdder;

public class OrderedThreadPool extends ThreadPoolExecutor {


    private final List<Executor> childExecutors = new ArrayList<>();
    private final Selector selector;

    public OrderedThreadPool(int corePoolSize,
                             int maximumPoolSize,
                             long keepAliveTime,
                             TimeUnit unit,
                             BlockingQueue<Runnable> workQueue,
                             ThreadFactory threadFactory,
                             RejectedExecutionHandler handler,
                             Selector selector) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        this.selector = selector;
        for (int i = 0; i < corePoolSize; i++) {
            ChildExecutor childExecutor = new ChildExecutor(workQueue.size() + workQueue.remainingCapacity());
            childExecutors.add(childExecutor);
        }
    }

    @Override
    public void execute(Runnable command) {
        if (command == null) {
            return;
        }
        if (command instanceof Ordered) {
            doOrderedExecute(command, ((Ordered) command).getKey());
        } else {
            doUnOrderedExecute(command);
        }
    }

    private void doOrderedExecute(Runnable command, String key) {
        childExecutors
                .get(selector.select(childExecutors.size(), key))
                .execute(command);
    }

    private void doUnOrderedExecute(Runnable command) {
        super.execute(command);
    }

    private final class ChildExecutor implements Executor, Runnable {
        private final BlockingQueue<Runnable> taskQueue;

        private final LongAdder completedTaskCount = new LongAdder();

        private final LongAdder rejectedTaskCount = new LongAdder();

        private boolean running;

        ChildExecutor(int queueSize) {
            taskQueue = new LinkedBlockingDeque<>(queueSize);
        }

        @Override
        public void execute(Runnable command) {
            boolean start = false;
            synchronized (this) {
                if (!taskQueue.add(command)) {
                    rejectedTaskCount.increment();
                    throw new RejectedExecutionException("task queue is full");
                }

                if (!running) {
                    start = true;
                    running = true;
                }
            }
            if (start) {
                doUnOrderedExecute(this);
            }
        }

        @Override
        public void run() {
            Runnable task;
            while (true) {
                try {
                    task = taskQueue.take();
                    try {
                        task.run();
                    } catch (RuntimeException x) {
                        throw x;
                    } finally {
                        completedTaskCount.increment();
                    }
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }
}
